package com.taotao.tools.core.concurrent.thread.pool;


import com.google.common.annotations.Beta;

import java.util.Collections;
import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

/**
 * <p>
 * 可以强制停止正在运行的任务。用于代码陷入死循环要杀死的场景
 * </p>
 *
 * @author fanbeibei
 * @version 1.0.0
 * @date 2023/11/30 17:04
 */
@Beta
public class StopableThreadPoolExecutor extends ThreadPoolExecutor {
    private final ConcurrentHashMap<Runnable, Thread> rtMap = new ConcurrentHashMap<>();
    private final Set<Runnable> stopRuns = Collections.synchronizedSet(new HashSet<>());


    public StopableThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue);
    }

    public StopableThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory);
    }

    public StopableThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, RejectedExecutionHandler handler) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, handler);
    }

    public StopableThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory, RejectedExecutionHandler handler) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory, handler);
    }


    @Override
    protected void beforeExecute(Thread t, Runnable r) {
        synchronized (this) {
            if (stopRuns.contains(r)) {
                throw new ThreadDeath();
            }

            rtMap.put(r, t);
        }
    }

    @Override
    protected void afterExecute(Runnable r, Throwable t) {
        rtMap.remove(r);
        stopRuns.remove(r);
    }


    /**
     * 强制停止执行的任务,会调用Thread.stop()方法。用于业务代码陷入死循序，无法interrupt退出场景
     *
     * @param command
     */
    public synchronized void forceStop(Runnable command) {
        stopRuns.add(command);
        Thread t = rtMap.get(command);
        if (null != t) {
            // 正在运行直接杀死
            t.stop();
            rtMap.remove(command);
            stopRuns.remove(command);
        }
    }
}
